package com.song.tools.songtest.config;

import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.aop.interceptor.AsyncUncaughtExceptionHandler;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.AsyncConfigurer;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Executor;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * SpringBoot 装配线程池,参考： <a href="https://blog.csdn.net/FBB360JAVA/article/details/120795737">SpringBoot使用线程池之ThreadPoolTaskExecutor和ThreadPoolExecutor</a>
 *
 * @author song tools
 * @since 2024-07-17
 */
@Slf4j
@EnableAsync
@Configuration
public class ThreadPoolConfig implements AsyncConfigurer {

    private static final String EXECUTOR_NAME = "asyncExecutor";

    @Value("${thread-pool.config.corePoolSize:10}")
    private Integer corePoolSize;
    @Value("${thread-pool.config.maxPoolSize:100}")
    private Integer maxPoolSize;
    @Value("${thread-pool.config.queueCapacity:200}")
    private Integer queueCapacity;
    @Value("${thread-pool.config.threadNamePrefix:AsyncThread-}")
    private String threadNamePrefix;
    @Value("${thread-pool.config.rejectedExecutionHandler:CallerRunsPolicy}")
    private String rejectedExecutionHandler;
    @Value("${thread-pool.config.preStartAllCoreThreads:false}")
    private Boolean preStartAllCoreThreads;

    private ThreadPoolTaskExecutor threadPoolTaskExecutor;

    @Bean(name = EXECUTOR_NAME)
    @Override
    public Executor getAsyncExecutor() {
        threadPoolTaskExecutor = new ThreadPoolTaskExecutor();
        // 核心线程数
        threadPoolTaskExecutor.setCorePoolSize(corePoolSize);
        // 最大线程数
        threadPoolTaskExecutor.setMaxPoolSize(maxPoolSize);
        // 阻塞队列容量
        threadPoolTaskExecutor.setQueueCapacity(queueCapacity);
        // 待任务在关机时完成--表明等待所有线程执行完
        threadPoolTaskExecutor.setWaitForTasksToCompleteOnShutdown(true);
        // 线程名称前缀
        threadPoolTaskExecutor.setThreadNamePrefix(threadNamePrefix);
        // 设置拒绝策略
        threadPoolTaskExecutor.setRejectedExecutionHandler(getRejectedExecutionHandler(rejectedExecutionHandler));
        // 核心线程预热
        threadPoolTaskExecutor.setPrestartAllCoreThreads(preStartAllCoreThreads);

        log.info("初始化线程池， corePoolSize={}, maxPoolSize={}, queueCapacity={}, threadNamePrefix={}, rejectedExecutionHandler={}, preStartAllCoreThreads={}", corePoolSize, maxPoolSize, queueCapacity, threadNamePrefix, rejectedExecutionHandler, preStartAllCoreThreads);
        threadPoolTaskExecutor.initialize();
        return threadPoolTaskExecutor;
    }


    @Override
    public AsyncUncaughtExceptionHandler getAsyncUncaughtExceptionHandler() {
        return (throwable, method, obj) -> {
            log.error("[ThreadPool Exception]：Message [{}], Method [{}]", throwable.getMessage(), method.getName());
            for (Object param : obj) {
                log.error("Parameter value [{}] ", param);
            }
        };
    }

    /**
     * 根据传入的参数获取拒绝策略
     *
     * @param rejectedName 拒绝策略名，比如 CallerRunsPolicy
     * @return RejectedExecutionHandler 实例对象，没有匹配的策略时，默认取 CallerRunsPolicy 实例
     */
    public RejectedExecutionHandler getRejectedExecutionHandler(String rejectedName) {
        Map<String, RejectedExecutionHandler> rejectedExecutionHandlerMap = new HashMap<>(16);
        rejectedExecutionHandlerMap.put("CallerRunsPolicy", new ThreadPoolExecutor.CallerRunsPolicy());
        rejectedExecutionHandlerMap.put("AbortPolicy", new ThreadPoolExecutor.AbortPolicy());
        rejectedExecutionHandlerMap.put("DiscardPolicy", new ThreadPoolExecutor.DiscardPolicy());
        rejectedExecutionHandlerMap.put("DiscardOldestPolicy", new ThreadPoolExecutor.DiscardOldestPolicy());
        return rejectedExecutionHandlerMap.getOrDefault(rejectedName, new ThreadPoolExecutor.CallerRunsPolicy());
    }

    @PreDestroy
    public void destroy() {
        if (threadPoolTaskExecutor != null) {
            ThreadPoolExecutor threadPoolExecutor = threadPoolTaskExecutor.getThreadPoolExecutor();
            threadPoolTaskExecutor.setWaitForTasksToCompleteOnShutdown(true);
            threadPoolExecutor.shutdown();
            try {
                if (!threadPoolExecutor.awaitTermination(1, TimeUnit.MINUTES)) {
                    threadPoolExecutor.shutdownNow();
                    if (!threadPoolExecutor.awaitTermination(5, java.util.concurrent.TimeUnit.SECONDS)){
                        log.info("ThreadPoolTaskExecutor shutdown fail");
                        return;
                    }
                }
            } catch (InterruptedException ex) {
                threadPoolExecutor.shutdownNow();
                Thread.currentThread().interrupt();
            }
            log.info("ThreadPoolTaskExecutor shutdown success");
        }
    }
}


